package main

import (
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"strings"
	"sync"
)

type ProjectAnalyzer struct {
	config    AnalyzerConfig
	fileSet   *token.FileSet
	packages  map[string]*ast.Package
	functions map[string]*FunctionInfo
	callGraph map[string][]string
	closures  map[string]*FunctionInfo // 存储闭包函数信息
	mu        sync.Mutex               // 添加互斥锁保证并发安全
}

type FunctionInfo struct {
	Name      string
	Package   string // 添加包信息
	File      string
	StartLine int
	EndLine   int
	Calls     []string
	CalledBy  []string
}

// 添加配置选项
type AnalyzerConfig struct {
	ExcludeDirs     []string // 要排除的目录
	ExcludePatterns []string // 要排除的文件模式
	MaxDepth        int      // 最大递归深度，-1表示无限制
	ShowExternal    bool     // 是否显示外部函数
	ShowFullPath    bool     // 是否显示完整文件路径
}

func NewProjectAnalyzer(config AnalyzerConfig) *ProjectAnalyzer {
	return &ProjectAnalyzer{
		config:    config,
		fileSet:   token.NewFileSet(),
		packages:  make(map[string]*ast.Package),
		functions: make(map[string]*FunctionInfo),
		callGraph: make(map[string][]string),
		closures:  make(map[string]*FunctionInfo),
	}
}

// 默认配置
func DefaultConfig() *AnalyzerConfig {
	return &AnalyzerConfig{
		ExcludeDirs: []string{
			"vendor",
			".git",
			"node_modules",
			"testdata",
		},
		ExcludePatterns: []string{
			"*_test.go",
			"*.pb.go",
		},
		MaxDepth:     -1,
		ShowExternal: false,
	}
}

func (pa *ProjectAnalyzer) AnalyzeDirectory(dir string, config *AnalyzerConfig) error {
	if config == nil {
		config = DefaultConfig()
	}

	// 获取项目根目录的绝对路径
	absRoot, err := filepath.Abs(dir)
	if err != nil {
		return err
	}

	var totalFiles, processedFiles int
	var fileList []string

	// 首先统计需要处理的文件总数
	err = filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// 检查是否需要跳过该目录
		if info.IsDir() {
			if shouldSkipDir(path, config.ExcludeDirs) {
				return filepath.SkipDir
			}

			// 检查递归深度
			if config.MaxDepth >= 0 {
				relPath, _ := filepath.Rel(absRoot, path)
				depth := len(strings.Split(relPath, string(os.PathSeparator)))
				if depth > config.MaxDepth {
					return filepath.SkipDir
				}
			}
			return nil
		}

		// 检查文件是否符合要求
		if strings.HasSuffix(path, ".go") && !shouldSkipFile(path, config.ExcludePatterns) {
			totalFiles++
			var targetPath string = path
			if config.ShowFullPath {
				abs, err := filepath.Abs(path)
				if err != nil {
					return err
				}
				targetPath = abs
			}
			fileList = append(fileList, targetPath)
		}
		return nil
	})

	if err != nil {
		return err
	}

	// 创建进度通道
	progress := make(chan struct{})
	go func() {
		for range progress {
			processedFiles++
			// fmt.Printf("\rAnalyzing files: %d/%d (%.2f%%)",
			// 	processedFiles, totalFiles,
			// 	float64(processedFiles)/float64(totalFiles)*100)
		}
	}()

	// 并发分析文件
	var wg sync.WaitGroup
	semaphore := make(chan struct{}, 10) // 限制并发数

	// 为每个go func生成唯一的名称
	closureCount := 0
	currentFunc := "main.AnalyzeDirectory"

	for _, file := range fileList {
		wg.Add(1)
		closureCount++
		gofuncName := fmt.Sprintf("%s.gofunc_%d", currentFunc, closureCount)

		go func(filename string, funcName string) {
			defer wg.Done()
			semaphore <- struct{}{} // 获取信号量
			defer func() {
				<-semaphore // 释放信号量
				progress <- struct{}{}
			}()

			if err := pa.analyzeFile(filename, absRoot, funcName); err != nil {
				fmt.Printf("\nError analyzing file %s: %v\n", filename, err)
			}
		}(file, gofuncName)
	}

	wg.Wait()
	close(progress)
	fmt.Println("\nAnalysis complete!")
	return nil
}

func (pa *ProjectAnalyzer) analyzeFile(filename, rootDir string, parentFunc string) error {
	pa.mu.Lock()
	defer pa.mu.Unlock()

	// 解析文件
	node, err := parser.ParseFile(pa.fileSet, filename, nil, parser.ParseComments)
	if err != nil {
		return err
	}

	// 获取相对于项目根目录的包路径
	relPath, _ := filepath.Rel(rootDir, filepath.Dir(filename))
	// 使用完整的目录路径作为包标识
	packagePath := strings.ReplaceAll(relPath, string(os.PathSeparator), "/")
	if packagePath == "." {
		packagePath = ""
	}

	// 获取包名
	packageName := node.Name.Name

	// 构建完整的包标识：目录路径 + 包名
	fullPackagePath := packageName
	if packagePath != "" {
		fullPackagePath = packagePath + "/" + packageName
	}

	fmt.Printf("Analyzing file: %s (package: %s)\n", filename, fullPackagePath)

	// 存储变量类型信息的映射
	varTypes := make(map[string]string)

	// 第一遍遍历，收集变量类型信息
	ast.Inspect(node, func(n ast.Node) bool {
		switch x := n.(type) {
		case *ast.AssignStmt:
			for i, lhs := range x.Lhs {
				if i >= len(x.Rhs) {
					break
				}
				if ident, ok := lhs.(*ast.Ident); ok {
					// 处理函数调用赋值
					if call, ok := x.Rhs[i].(*ast.CallExpr); ok {
						if funcIdent, ok := call.Fun.(*ast.Ident); ok {
							// 从函数名推断类型
							// 例如：NewXXX() 返回 XXX 类型
							typeName := funcIdent.Name
							if strings.HasPrefix(typeName, "New") {
								typeName = strings.TrimPrefix(typeName, "New")
							}
							varTypes[ident.Name] = typeName
						}
					}
				}
			}
		}
		return true
	})

	// 闭包计数器
	closureCount := 0

	// 当前正在分析的函数名
	currentFunc := parentFunc

	// 第二遍遍历，分析函数和调用关系
	ast.Inspect(node, func(n ast.Node) bool {
		switch x := n.(type) {
		case *ast.GoStmt:
			// 处理 go func() 语句
			if funcLit, ok := x.Call.Fun.(*ast.FuncLit); ok {
				closureCount++
				closureName := fmt.Sprintf("%s.anonymous_%d", currentFunc, pa.fileSet.Position(x.Pos()).Line)

				position := pa.fileSet.Position(funcLit.Pos())
				endPosition := pa.fileSet.Position(funcLit.End())

				pa.closures[closureName] = &FunctionInfo{
					Name:      closureName,
					Package:   fullPackagePath,
					File:      filename,
					StartLine: position.Line,
					EndLine:   endPosition.Line,
					Calls:     make([]string, 0),
					CalledBy:  make([]string, 0),
				}

				// 保存当前函数名以分析闭包内的调用
				previousFunc := currentFunc
				currentFunc = closureName

				// 分析闭包内的函数调用
				ast.Inspect(funcLit.Body, func(n ast.Node) bool {
					if call, ok := n.(*ast.CallExpr); ok {
						pa.analyzeCall(call, closureName, fullPackagePath, varTypes, filename)
					}
					return true
				})

				// 恢复当前函数名
				currentFunc = previousFunc
			}

		case *ast.FuncLit:
			// 处理普通闭包函数
			closureCount++
			closureName := fmt.Sprintf("%s.anonymous_%d", currentFunc, pa.fileSet.Position(x.Pos()).Line)

			position := pa.fileSet.Position(x.Pos())
			endPosition := pa.fileSet.Position(x.End())

			pa.closures[closureName] = &FunctionInfo{
				Name:      closureName,
				Package:   fullPackagePath,
				File:      filename,
				StartLine: position.Line,
				EndLine:   endPosition.Line,
				Calls:     make([]string, 0),
				CalledBy:  make([]string, 0),
			}

			// 保存当前函数名以分析闭包内的调用
			previousFunc := currentFunc
			currentFunc = closureName

			// 分析闭包内的函数调用
			ast.Inspect(x.Body, func(n ast.Node) bool {
				if call, ok := n.(*ast.CallExpr); ok {
					pa.analyzeCall(call, closureName, fullPackagePath, varTypes, filename)
				}
				return true
			})

			// 恢复当前函数名
			currentFunc = previousFunc

		case *ast.FuncDecl:
			funcName := x.Name.Name
			if x.Recv != nil {
				if recv, ok := x.Recv.List[0].Type.(*ast.StarExpr); ok {
					if ident, ok := recv.X.(*ast.Ident); ok {
						funcName = ident.Name + "." + funcName
					}
				} else if ident, ok := x.Recv.List[0].Type.(*ast.Ident); ok {
					funcName = ident.Name + "." + funcName
				}
			}

			// 添加包名前缀
			currentFunc = fullPackagePath + "." + funcName
			fmt.Printf("Found function: %s\n", currentFunc)

			position := pa.fileSet.Position(x.Pos())
			endPosition := pa.fileSet.Position(x.End())

			pa.functions[currentFunc] = &FunctionInfo{
				Name:      funcName,
				Package:   fullPackagePath,
				File:      filename,
				StartLine: position.Line,
				EndLine:   endPosition.Line,
				Calls:     make([]string, 0),
				CalledBy:  make([]string, 0),
			}

		case *ast.CallExpr:
			if currentFunc != "" {
				pa.analyzeCall(x, currentFunc, fullPackagePath, varTypes, filename)
			}
		}
		return true
	})

	return nil
}

// 抽取函数调用分析逻辑到单独的方法
func (pa *ProjectAnalyzer) analyzeCall(x *ast.CallExpr, currentFunc, packagePath string, varTypes map[string]string, filename string) {
	var calleeName string
	callPos := pa.fileSet.Position(x.Pos()) // 获取调用位置

	switch fun := x.Fun.(type) {
	case *ast.SelectorExpr:
		if ident, ok := fun.X.(*ast.Ident); ok {
			if ident.Obj == nil {
				// 外部包函数调用
				calleeName = ident.Name + "." + fun.Sel.Name
			} else {
				// 本地方法调用
				if typeName, exists := varTypes[ident.Name]; exists {
					calleeName = packagePath + "." + typeName + "." + fun.Sel.Name
				} else if field, ok := ident.Obj.Decl.(*ast.Field); ok {
					if starExpr, ok := field.Type.(*ast.StarExpr); ok {
						if typeIdent, ok := starExpr.X.(*ast.Ident); ok {
							calleeName = packagePath + "." + typeIdent.Name + "." + fun.Sel.Name
						}
					} else if typeIdent, ok := field.Type.(*ast.Ident); ok {
						calleeName = packagePath + "." + typeIdent.Name + "." + fun.Sel.Name
					}
				}
			}
		}
	case *ast.Ident:
		calleeName = packagePath + "." + fun.Name
	case *ast.FuncLit:
		// 直接调用的匿名函数
		closureName := fmt.Sprintf("%s.anonymous_%d", currentFunc, callPos.Line)
		calleeName = closureName

		// 保存匿名函数信息
		endPosition := pa.fileSet.Position(fun.End())
		pa.closures[closureName] = &FunctionInfo{
			Name:      closureName,
			Package:   packagePath,
			File:      filename,
			StartLine: callPos.Line, // 使用调用位置
			EndLine:   endPosition.Line,
			Calls:     make([]string, 0),
			CalledBy:  make([]string, 0),
		}
	}

	if calleeName != "" {
		if _, exists := pa.callGraph[currentFunc]; !exists {
			pa.callGraph[currentFunc] = make([]string, 0)
		}

		// 检查是否已经添加过这个调用关系
		found := false
		for _, existing := range pa.callGraph[currentFunc] {
			if existing == calleeName {
				found = true
				break
			}
		}
		if !found {
			// 保存调用位置信息
			callInfo := fmt.Sprintf("%s (%s:%d)", calleeName, filename, callPos.Line)
			pa.callGraph[currentFunc] = append(pa.callGraph[currentFunc], callInfo)
		}
	}
}

// AnalyzeFromEntry 从指定的入口函数开始分析调用关系
func (pa *ProjectAnalyzer) AnalyzeFromEntry(entryFunc string) {
	fmt.Printf("\nStarting analysis from entry point: %s\n", entryFunc)

	visited := make(map[string]bool)
	var dfs func(funcName string, depth int, isLast bool, prefix string)

	dfs = func(funcName string, depth int, isLast bool, prefix string) {
		if visited[funcName] {
			return
		}
		visited[funcName] = true

		// 构建树形显示的前缀
		var currentPrefix string
		if depth > 0 {
			if isLast {
				currentPrefix = prefix + "└── "
			} else {
				currentPrefix = prefix + "├── "
			}
		}

		// 从函数名中提取实际名称和调用位置
		actualName := funcName
		callLocation := ""
		if idx := strings.LastIndex(funcName, " ("); idx != -1 {
			actualName = funcName[:idx]
			callLocation = funcName[idx:]
		}

		// 获取函数信息
		info, exists := pa.functions[actualName]
		if !exists {
			info, exists = pa.closures[actualName]
		}

		// 打印当前函数
		if exists {
			if callLocation != "" {
				fmt.Printf("%s%s%s\n", currentPrefix, actualName, callLocation)
			} else {
				fmt.Printf("%s%s (%s:%d)\n", currentPrefix, actualName, info.File, info.StartLine)
			}

			// 获取所有被调用的函数
			callees := pa.callGraph[actualName]

			// 计算下一层的前缀
			var nextPrefix string
			if depth > 0 {
				if isLast {
					nextPrefix = prefix + "    "
				} else {
					nextPrefix = prefix + "│   "
				}
			}

			// 递归分析该函数调用的其他函数
			for i, callee := range callees {
				isLastCallee := i == len(callees)-1
				// 检查是否为外部函数
				actualCallee := callee
				if idx := strings.LastIndex(callee, " ("); idx != -1 {
					actualCallee = callee[:idx]
				}
				_, isInternal := pa.functions[actualCallee]
				_, isClosure := pa.closures[actualCallee]
				if isInternal || isClosure || pa.config.ShowExternal {
					dfs(callee, depth+1, isLastCallee, nextPrefix)
				}
			}
		} else if pa.config.ShowExternal {
			fmt.Printf("%s%s%s\n", currentPrefix, actualName, callLocation)
		}
	}

	fmt.Printf("\nCall tree from entry point: %s\n", entryFunc)
	fmt.Println("─────────────────────────────────────")
	dfs(entryFunc, 0, true, "")
}

func shouldSkipDir(path string, excludeDirs []string) bool {
	base := filepath.Base(path)
	for _, exclude := range excludeDirs {
		if base == exclude {
			return true
		}
	}
	return false
}

func shouldSkipFile(path string, patterns []string) bool {
	for _, pattern := range patterns {
		matched, _ := filepath.Match(pattern, filepath.Base(path))
		if matched {
			return true
		}
	}
	return false
}

// PrintAnalysis 增加了更详细的输出
func (pa *ProjectAnalyzer) PrintAnalysis() {
	fmt.Println("\n=== Project Analysis Results ===")

	// 按包组织函数
	packageFuncs := make(map[string][]*FunctionInfo)
	for _, info := range pa.functions {
		packageFuncs[info.Package] = append(packageFuncs[info.Package], info)
	}

	// 打印包级别的信息
	for pkg, funcs := range packageFuncs {
		fmt.Printf("\nPackage: %s\n", pkg)
		fmt.Printf("Number of functions: %d\n", len(funcs))

		for _, info := range funcs {
			fmt.Printf("\n  Function: %s\n", info.Name)
			fmt.Printf("    File: %s\n", info.File)
			fmt.Printf("    Lines: %d-%d\n", info.StartLine, info.EndLine)

			if len(info.CalledBy) > 0 {
				fmt.Printf("    Called By:\n")
				for _, caller := range info.CalledBy {
					fmt.Printf("      - %s\n", caller)
				}
			}
		}
	}

	// 打印调用图
	fmt.Println("\nCall Graph:")
	for caller, callees := range pa.callGraph {
		if len(callees) > 0 {
			fmt.Printf("\n%s calls:\n", caller)
			for _, callee := range callees {
				fmt.Printf("  -> %s\n", callee)
			}
		}
	}
}
